package com.soup.lab.basic.iso;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

public class MultiKeyTaskManager<K> {

    private Map<K, NodeLink> table = new ConcurrentHashMap<>();

    public void addTask(Runnable runnable, Collection<K> keys) {
        keys.forEach(
                key -> table.compute(key, (k, v) -> {
                    NodeLink link = v;
                    if (v == null) {
                        link = new NodeLink();
                    }
                    link.linkTail(new Node(runnable));
                    return link;
                }));
    }

    public Runnable getTask(K key) {
        NodeLink link = table.get(key);
        Node node;
        if(link == null || (node = link.unlinkHead()) == null) {
            return null;
        }
        return node;
    }

    private static class NodeLink {
        Node head;
        Node tail;
        Object monitor = new Object();

        void linkTail(Runnable r) {
            synchronized (monitor) {
                Node n = new Node(r);
                if (tail == null) {
                    head = tail = n;
                } else {
                    tail.next = n;
                    n.waiting(tail);
                    tail = n;
                }
            }
        }

        Node unlinkHead() {
            synchronized (monitor) {
                Node h = head;
                if (h != null) {
                    head = h.next;
                }
                return h;
            }
        }

        @Override
        public String toString() {
            synchronized (monitor) {
                StringBuilder sb = new StringBuilder();
                Node node = head;
                while(node != null) {
                    sb.append(node.toString()).append(System.lineSeparator());
                    node = node.next;
                }
                return sb.toString();
            }
        }
    }


    private static class Node implements Runnable {
        final Runnable runnable;
        // 后继节点
        volatile Node next;
        // 当前节点依赖的前驱节点
        final Set<Node> prevSet;
        final Object prevSetMonitor = new Object();

        /**
         * 节点的状态
         * -1-waiting:
         * 等待中，当前节点正在等待其他节点执行完成
         * 0-runnable:
         * 可执行，当前节点不依赖其他节点，可以直接执行
         * 1-running
         * 正在执行，当前节点正在执行
         * 2-finished
         * 已完成，当前节点执行完成
         */
        final AtomicInteger state = new AtomicInteger(RUNNABLE);
        static final int WAITING = -1;
        static final int RUNNABLE = 0;
        static final int RUNNING = 1;
        static final int FINISHED = 2;

        public Node(Runnable runnable) {
            this.runnable = runnable;
            this.prevSet = new LinkedHashSet<>();
        }

        public Node(Runnable runnable, Collection<Node> prevSet) {
            this.runnable = runnable;
            this.prevSet = new LinkedHashSet<>(prevSet);
        }

        @Override
        public void run() {
            for (; ; ) {
                synchronized (prevSetMonitor) {
                    if (!state.compareAndSet(RUNNABLE, RUNNING)) {
                        checkWaiting();
                        continue;
                    }
                }
                try {
                    runnable.run();
                } finally {
                    finish();
                    break;
                }
            }
        }

        /**
         * 当前节点执行完成后，应通知后面等待的节点，调用它的 finishWaiting 方法
         */
        void finish() {
            for (; ; ) {
                int status = state.get();
                if (status != RUNNING) {
                    break;
                }
                if (state.compareAndSet(RUNNING, FINISHED) && next != null) {
                    next.finishWaiting(this);
                    break;
                }
            }
        }

        /**
         * 当前节点执行完成
         *
         * @param prev
         */
        void finishWaiting(Node prev) {
            synchronized (prevSetMonitor) {
                prevSet.remove(prev);
                if (prevSet.isEmpty()) {
                    state.compareAndSet(WAITING, RUNNABLE);
                }
            }
        }

        void waiting(Node prev) {
            synchronized (prevSetMonitor) {
                if (prev.state.get() > RUNNING && state.get() <= RUNNABLE) {
                    prevSet.add(prev);
                    for (; ; ) {
                        if (state.get() == WAITING) {
                            break;
                        }
                        state.compareAndSet(RUNNABLE, WAITING);
                    }
                }
            }
        }

        void checkWaiting() {
            synchronized (prevSetMonitor) {
                Iterator<Node> iterator = prevSet.iterator();
                while (iterator.hasNext()) {
                    if (iterator.next().getState() != FINISHED) {
                        return;
                    }
                    iterator.remove();
                }
                state.compareAndSet(WAITING, RUNNABLE);
            }
        }

        int getState() {
            return state.get();
        }

        @Override
        public String toString() {
            synchronized (prevSetMonitor) {
                return "状态：" + state.get() + ", 前驱节点："+ prevSet;
            }
        }
    }

}
